﻿Imports System.Runtime.CompilerServices
Imports System.Text
Imports System.Diagnostics

Namespace Extension
    Public Module VisualTreeHelperExtensions
        <Extension>
        Public Function FindVisualChild(root As FrameworkElement, name As String) As FrameworkElement
            Dim temp = DirectCast(root.FindName(name), FrameworkElement)
            If temp IsNot Nothing Then
                Return temp
            End If

            For Each e In root.GetVisualChildren
                temp = DirectCast(e.FindName(name), FrameworkElement)
                If temp IsNot Nothing Then
                    Return temp
                End If
            Next

            Return Nothing
        End Function

        <Extension>
        Public Function GetVisualParent(node As FrameworkElement) As FrameworkElement
            Return DirectCast(VisualTreeHelper.GetParent(node), FrameworkElement)
        End Function

        <Extension>
        Public Function GetVisualChild(node As FrameworkElement, index As Integer) As FrameworkElement
            Return DirectCast(VisualTreeHelper.GetChild(node, index), FrameworkElement)
        End Function

        <Extension>
        Public Iterator Function GetVisualChildren(root As FrameworkElement) As IEnumerable(Of FrameworkElement)
            For i = 0 To VisualTreeHelper.GetChildrenCount(root) - 1
                Yield root.GetVisualChild(i)
            Next
        End Function

        <Extension>
        Public Iterator Function GetVisualAncestors(node As FrameworkElement) As IEnumerable(Of FrameworkElement)
            Dim p = node.GetVisualParent
            While p IsNot Nothing
                Yield p
                p = p.GetVisualParent
            End While
        End Function

        <Extension>
        Public Function GetVisualStateGroup(root As FrameworkElement, groupName As String, searchAncestors As Boolean) As VisualStateGroup
            Dim groups = VisualStateManager.GetVisualStateGroups(root)
            For Each o In groups
                Dim group = DirectCast(o, VisualStateGroup)
                If group IsNot Nothing AndAlso group.Name = groupName Then
                    Return group
                End If
            Next

            If searchAncestors Then
                Dim p = root.GetVisualParent
                If p IsNot Nothing Then
                    Return p.GetVisualStateGroup(groupName, True)
                End If
            End If

            Return Nothing
        End Function

        <Extension>
        Public Function FindVisualState(root As FrameworkElement, name As String) As VisualStateGroup
            If root Is Nothing Then
                Return Nothing
            End If

            Dim groups = VisualStateManager.GetVisualStateGroups(root)
            Return groups.Cast(Of VisualStateGroup).FirstOrDefault(Function(g) g.Name = name)
        End Function

        <Extension>
        Public Iterator Function GetVisualDescendents(root As FrameworkElement) As IEnumerable(Of FrameworkElement)
            Dim toDo As New Queue(Of IEnumerable(Of FrameworkElement))

            toDo.Enqueue(root.GetVisualChildren)
            While toDo.Count > 0
                Dim cs = toDo.Dequeue
                For Each c In cs
                    Yield c
                    toDo.Enqueue(c.GetVisualChildren)
                Next
            End While
        End Function

        <Extension>
        <Conditional("DEBUG")>
        Public Sub GetVisualChildTreeDebugText(root As FrameworkElement, result As StringBuilder)
            Dim rs As New List(Of String)
            root.GetChildTree("", "  ", rs)
            For Each s In rs
                result.AppendLine(s)
            Next
        End Sub

        <Extension>
        Public Sub GetChildTree(root As FrameworkElement, prefix As String, addPrefix As String, results As List(Of String))
            Dim thisElement = ""
            If String.IsNullOrEmpty(root.Name) Then
                thisElement = "[Anonymous]"
            Else
                thisElement = String.Format("[{0}]", root.Name)
            End If

            thisElement &= String.Format(" : {0}", root.GetType.Name)

            results.Add(prefix & thisElement)
            For Each c In root.GetVisualChildren
                c.GetChildTree(prefix & addPrefix, addPrefix, results)
            Next
        End Sub

        <Extension>
        <Conditional("DEBUG")>
        Public Sub GetAncestorVisualTreeDebugText(node As FrameworkElement, result As StringBuilder)
            Dim tree As New List(Of String)
            node.GetAncestorVisualTree(tree)
            Dim prefix = ""
            For Each s In tree
                result.AppendLine(prefix & s)
                prefix &= "  "
            Next
        End Sub

        <Extension>
        Public Sub GetAncestorVisualTree(node As FrameworkElement, children As List(Of String))
            Dim name = If(String.IsNullOrEmpty(node.Name), "[Anon]", node.Name)
            Dim thisNode = name & ": " & node.GetType.Name

            children.Insert(0, thisNode)
            Dim p = node.GetVisualParent
            If p IsNot Nothing Then
                GetAncestorVisualTree(p, children)
            End If
        End Sub

        <Extension>
        Public Function GetTransform(Of TRequestedTransform As {Transform, New})(element As UIElement, mode As TransformCreationMode) As TRequestedTransform
            Dim originalTransform As Transform = element.RenderTransform
            Dim requestedTransform As TRequestedTransform = Nothing
            Dim matrixTransform As MatrixTransform = Nothing
            Dim transformGroup As TransformGroup = Nothing

            ' Current transform is null -- create if necessary and return
            If originalTransform Is Nothing Then
                If (mode And TransformCreationMode.Create) = TransformCreationMode.Create Then
                    requestedTransform = New TRequestedTransform()
                    element.RenderTransform = requestedTransform
                    Return requestedTransform
                End If

                Return Nothing
            End If

            ' Transform is exactly what we want -- return it
            requestedTransform = TryCast(originalTransform, TRequestedTransform)
            If requestedTransform IsNot Nothing Then
                Return requestedTransform
            End If


            ' The existing transform is matrix transform - overwrite if necessary and return
            matrixTransform = TryCast(originalTransform, MatrixTransform)
            If matrixTransform IsNot Nothing Then
                If matrixTransform.Matrix.IsIdentity AndAlso (mode And TransformCreationMode.Create) = TransformCreationMode.Create AndAlso (mode And TransformCreationMode.IgnoreIdentityMatrix) = TransformCreationMode.IgnoreIdentityMatrix Then
                    requestedTransform = New TRequestedTransform()
                    element.RenderTransform = requestedTransform
                    Return requestedTransform
                End If

                Return Nothing
            End If

            ' Transform is actually a group -- check for the requested type
            transformGroup = TryCast(originalTransform, TransformGroup)
            If transformGroup IsNot Nothing Then
                For Each child As Transform In transformGroup.Children
                    ' Child is the right type -- return it
                    If TypeOf child Is TRequestedTransform Then
                        Return TryCast(child, TRequestedTransform)
                    End If
                Next

                ' Right type was not found, but we are OK to add it
                If (mode And TransformCreationMode.AddToGroup) = TransformCreationMode.AddToGroup Then
                    requestedTransform = New TRequestedTransform()
                    transformGroup.Children.Add(requestedTransform)
                    Return requestedTransform
                End If

                Return Nothing
            End If

            ' Current ransform is not a group and is not what we want;
            ' create a new group containing the existing transform and the new one
            If (mode And TransformCreationMode.CombineIntoGroup) = TransformCreationMode.CombineIntoGroup Then
                transformGroup = New TransformGroup()
                transformGroup.Children.Add(originalTransform)
                transformGroup.Children.Add(requestedTransform)
                element.RenderTransform = transformGroup
                Return requestedTransform
            End If

            Debug.Assert(False, "Shouldn't get here")
            Return Nothing
        End Function

        <Extension>
        Public Function GetTransformPropertyPath(Of TRequestedType As Transform)(element As FrameworkElement, subProperty As String) As String
            Dim t As Transform = element.RenderTransform
            If TypeOf t Is TRequestedType Then
                Return String.Format("(RenderTransform).({0}.{1})", GetType(TRequestedType).Name, subProperty)

            ElseIf TypeOf t Is TransformGroup Then
                Dim g As TransformGroup = TryCast(t, TransformGroup)
                For i As Integer = 0 To g.Children.Count - 1
                    If TypeOf g.Children(i) Is TRequestedType Then
                        Return String.Format("(RenderTransform).(TransformGroup.Children)[" & i & "].({0}.{1})", GetType(TRequestedType).Name, subProperty)
                    End If
                Next
            End If

            Return ""
        End Function

        <Extension>
        Public Function GetPlaneProjection(element As UIElement, create As Boolean) As PlaneProjection
            Dim originalProjection = element.Projection
            Dim p As PlaneProjection = Nothing

            If TypeOf originalProjection Is PlaneProjection Then
                Return DirectCast(originalProjection, PlaneProjection)
            End If

            If originalProjection Is Nothing Then
                If create Then
                    p = New PlaneProjection
                    element.Projection = p
                End If
            End If

            Return p
        End Function
    End Module

    <Flags>
    Public Enum TransformCreationMode
        ''' <summary>
        ''' Don't try and create a transform if it doesn't already exist
        ''' </summary>
        None = 0

        ''' <summary>
        ''' Create a transform if none exists
        ''' </summary>
        Create = 1

        ''' <summary>
        ''' Create and add to an existing group
        ''' </summary>
        AddToGroup = 2

        ''' <summary>
        ''' Create a group and combine with existing transform; may break existing animations
        ''' </summary>
        CombineIntoGroup = 4

        ''' <summary>
        ''' Treat identity matrix as if it wasn't there; may break existing animations
        ''' </summary>
        IgnoreIdentityMatrix = 8

        ''' <summary>
        ''' Create a new transform or add to group
        ''' </summary>
        CreateOrAddAndIgnoreMatrix = Create Or AddToGroup Or IgnoreIdentityMatrix

        ''' <summary>
        ''' Default behaviour, equivalent to CreateOrAddAndIgnoreMatrix
        ''' </summary>
        [Default] = CreateOrAddAndIgnoreMatrix
    End Enum
End Namespace